% EKF based localization of a robot with unknown landmarks, with known
% identities

clear; clc; close all;
%load Q5_example_data.mat
load Q5_data.mat

T = length(u);
ss = length(x0);
xfilt = x0;
Sigma_filt{1} = Sigma0;

landmarks_discovered = zeros(num_landmarks,1);
num_landmarks_discovered = 0;

for t=1:T-1

    % your code for dynamics update here
    xfilt(:,t+1) = f_robot(xfilt(:,t), u(:,t), dt);
    A = jacobian_f_robot(xfilt(:,t), u(:,t), dt);
	Sigma_filt{t+1} = A*Sigma_filt{t}*A' + Q*dt;% your code here
    
    for l=1:num_landmarks
        if(y_landmarks_valid{l}(t+1))
            if(landmarks_discovered(l))
                % if discovered in the past, the landmark position is
                % already in filter state, and this is a EKF measurement
                % update:
                

                ypred = f_unknown_landmark(xfilt(:,t+1), landmark_discovery_idx(l)); % your code into f_unknown_landmark.m
                y = y_landmarks{l}(:,t+1);

                C = jacobian_f_unknown_landmark(xfilt(:,t+1), landmark_discovery_idx(l)); % your code into jacobian_f_unknown_landmark.m

                % your code here
                K = Sigma_filt{t+1}*C'*inv(C*Sigma_filt{t+1}*C' +  R_landmark);
                
                xfilt(:,t+1) = xfilt(:,t+1) + K * (y - ypred);% your code here
                Sigma_filt{t+1} = Sigma_filt{t+1} - K*C*Sigma_filt{t+1};% your code here
            else
                % add into filter state:

                % bookkeeping on the landmarks discovered:
                num_landmarks_discovered = num_landmarks_discovered + 1;
                landmark_discovery_idx(l) = num_landmarks_discovered;
                landmark_discovery_time(l) = t+1;
                landmarks_discovered(l) = 1;

                % statesize increases by 2 (the two coordinates of the
                % landmark
                ss = ss + 2;

                % augment the state
                % padd zeros for past times for the landmark
                xfilt = [xfilt; zeros(2,size(xfilt,2))];
                
                % compute an initial estimate of landmark north,east position:
                R = [cos(xfilt(3,t+1)) -sin(xfilt(3,t+1)); sin(xfilt(3,t+1)) cos(xfilt(3,t+1))];
                %rotation matrix between robot and world frame:
                landmark_robot = y_landmarks{l}(:,t+1);
                
                
                landmark_ne_estimate = R*landmark_robot + [xfilt(1,t+1) ; xfilt(2,t+1)] ; % your code here
                xfilt(end-1:end,t+1) = landmark_ne_estimate;
                
                % augment process noise matrix (no noise on
                % landmarks---they are stationary)
                Q = [ Q  zeros(size(Q,1),2) ; zeros(2, size(Q,2)+2)]; % no process noise on the landmarks

                % augment filter covariance -- this is somewhat a design choice, we chose the following:
                Sigma_filt{t+1} = [Sigma_filt{t+1} zeros(size(Sigma_filt{t+1},1),2); zeros(2, size(Sigma_filt{t+1},2)) zeros(2,2)];
                C = jacobian_f_unknown_landmark(xfilt(:,t+1), num_landmarks_discovered);
                landmark_ne_cov = R' * (C*Sigma_filt{t+1}*C' + 6*R_landmark) * R; %Innovation (or residual) covariance
                Sigma_filt{t+1}(end-1:end,end-1:end) = landmark_ne_cov;
            end
        end
    end
end        

% plot subsampled trajectory:
colors1 = ['km']; colors2 = ['bg'];
map_fig_id = figure; hold on; axis([-5 25 -5 25]); axis equal; xlabel('East'); ylabel('North');
spacing = 20;
for i= [1:spacing:size(xfilt,2) size(xfilt,2)]
    plot_uncertainty_ellipse(xfilt(1:3,i), Sigma_filt{i}, map_fig_id, colors1);
    for l=1:num_landmarks_discovered
       if(landmark_discovery_time(l) < i)
           l_state_idxs = 3+ 2*(landmark_discovery_idx(l)-1)+1: 3 + 2*(landmark_discovery_idx(l)-1)+2;
           plot_uncertainty_ellipse(xfilt(l_state_idxs,i), Sigma_filt{i}(l_state_idxs,l_state_idxs), map_fig_id, colors2);
       end
    end
end

% plot final map:
final_map_fig_id = figure; hold on; axis([-5 25 -5 25]); axis equal; xlabel('East'); ylabel('North'); title('map');
colors3 = ['kr'];
for l=1:num_landmarks_discovered
    l_state_idxs = 3+ 2*(l-1)+1: 3 + 2*(l-1)+2;
    %if(trace(Sigma_filt(l_state_idxs,l_state_idxs,i)) < 2*sigma0_landmarks^2 / 10)
        %% if we have seen the landmark a few times the variance will have decreased and we plot it
        plot_uncertainty_ellipse(xfilt(l_state_idxs,end), Sigma_filt{end}(l_state_idxs,l_state_idxs), final_map_fig_id, colors3);
    %end
end
axis equal;



